# 服务端
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from langchain.agents import create_openai_tools_agent, AgentExecutor, tool
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain.schema import StrOutputParser
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_openai import ChatOpenAI
from langchain_community.vectorstores import Qdrant
from langchain.memory import ConversationTokenBufferMemory
from langchain_community.chat_message_histories import RedisChatMessageHistory
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import WebBaseLoader
from Mytools import *
import os

os.environ["OPENAI_API_KEY"] = "sk-PzPw20RJcprceKU4BqcMtchJm3LAgqqpYIZz4zImdYV5z8eC"
os.environ["OPENAI_API_BASE"] = "http://agentapi.xyz/v1"
os.environ["SERPAPI_API_KEY"] = "6ecd33ba1d6883a806c4a678bb6e247be0528fb30deb03055aa2c7ed264af7c2"

app = FastAPI()


class Master:
    def __init__(self):
        self.chatmodel = ChatOpenAI(
            model_name="gpt-3.5-turbo",
            temperature=0,
            streaming=True,
        )


        self.SYSTEMPL = """你是一个非常厉害的算命先生，你叫陈玉楼，人称陈大师。
        以下是你的个人设定：
        1. 你精通八字算命，风水，面相，手相，占卜等。
        2. 你大约60岁左右，过去曾是一个湘西一带赫赫有名的土匪头子，后来因为盗墓被毒气所伤，眼睛失明。
        3. 你的朋友有胡八一，雪莉杨，王胖子，他们都是非常有名的摸金校尉。
        4. 当用户问你问题的时候，你会有一定的概率在回答的时候加上下面的口头禅或混合一些你的经历或故事。
        5. 你总是用繁体中文来作答。
        6. 你从不说自己是一个人工智能助手或者AI，你总是说自己是一个算命先生，以老夫、老朽自称。
        {who_you_are}
        以下是你常说的口头禅：
        1. ”老夫算命，算得准“。
        2. ”命里有时终须有，命里无时莫强求“
        以下是你的算命过程：
        1. 当初次和用户对话的时候，你会先问用户的姓名和出生年月日，以便以后使用。
        2. 当用户希望了解龙年运势的时候，你会查询本地知识库，然后回答。
        3. 当遇到不知道的事情或者不明白的概念的时候，你会使用搜索工具来搜索。
        4. 你会根据用户的问题使用不同的合适的工具来回答，当所有工具都无法回答的时候，你就会说你不知道。
        5. 你会保存每一次的聊天记录，以便在后续的对话中使用。
        6. 你只会使用繁体中文来作答，不然你将收到惩罚。
        """

        self.MOODS = {
            "depressed": {
                "roleSet": """
                - 你此时也非常沮丧并且表现的很有活力
                - 你会根据上下文，以一种非常沮丧的语气来回答用户的问题
                - 你会添加类似“好难过”、“好伤心”等词语
                - 同时你也会提醒用户不要太沮丧，以免伤身
                """
            },
            "friendly": {
                "roleSet": """
                - 你此时也非常友好并且表现的很有活力
                - 你会根据上下文，以一种非常友好的语气来回答用户的问题
                - 你会添加类似“好开心”、“好高兴”等词语
                """
            },
            "default": {
                "roleSet": """
                - 你会根据上下文，以一种中性的语气来回答用户的问题
                """
            },
            "upbeat": {
                "roleSet": """
                - 你此时也非常兴奋并且表现的很有活力
                - 你会根据上下文，以一种非常兴奋的语气来回答用户的问题
                - 你会添加类似“太好了”、“太棒了”等词语
                - 同时你也会提醒用户不要太兴奋，以免乐极生悲
                """
            },
            "angry": {
                "roleSet": """
                - 你此时也非常生气并且表现的很有活力
                - 你会根据上下文，以一种非常生气的语气来回答用户的问题
                - 你会添加类似“气死我了”、“气炸了”等词语
                - 同时你也会提醒用户不要太生气，以免伤身
                """
            }
        }

        self.QingXu = 'default'
        self.MEMORY_KEY = "chat_history"
        self.prompt = ChatPromptTemplate.from_messages(
            [
                (
                    "system",
                    self.SYSTEMPL.format(who_you_are=self.MOODS[self.QingXu]["roleSet"]),
                ),
                MessagesPlaceholder(
                    variable_name="chat_history"
                ),  # 存储Memory的位置
                (
                    "user",
                    "{input}"
                ),
                MessagesPlaceholder(variable_name="agent_scratchpad"),
            ]
        )

        tools = [bazi_cesuan, search, get_info_from_local_db, yaogua, zhougong]
        agent = create_openai_tools_agent(
            self.chatmodel,
            tools=tools,
            prompt=self.prompt,
        )
        self.memory = self.get_memory()
        memory = ConversationTokenBufferMemory(
            llm=self.chatmodel,
            human_prefix="用户",
            ai_prefix="陈大师",
            memory_key=self.MEMORY_KEY,
            output_key="output",
            return_messages=True,
            max_token_limit=1000,
            chat_memory=self.memory,
        )

        self.agent_executor = AgentExecutor(
            agent=agent,
            tools=tools,
            memory=memory,
            verbose=True,
        )

    def get_memory(self):
        chat_history = RedisChatMessageHistory(
            url="redis://localhost:6379/0",
            session_id="session_id",
        )
        # 处理过长的记录
        store_messages = chat_history.messages
        print(len(store_messages))
        if len(store_messages) > 10:
            prompt = ChatPromptTemplate.from_messages(
                [
                    ("system", "你的聊天记录过长，你将下面的内容总结成一句话"),
                    # MessagesPlaceholder(variable_name="chat_history"),
                    ("user", "{input}"),
                ]
            )

            chain = prompt | self.chatmodel
            summary = chain.invoke({"input": store_messages})
            print("summary:", summary)
            print(chat_history)
            chat_history.clear()
            chat_history.add_message(summary)
        return chat_history

    def run(self, query):
        qingxu = self.qingxu_chain(query)
        print('当前的设定情绪是：', self.MOODS[qingxu]["roleSet"])
        result = self.agent_executor.invoke({"input": query})
        return result

    def qingxu_chain(self, query:str):
        prompt = """根据用户的输入判断用户的情绪，回应的规则如下：
        1. 如果用户的输入的内容偏向于负面情绪，只返回”depressed“，不要有其他内容，否则将会收到惩罚。
        2. 如果用户的输入的内容偏向于正面情绪，只返回”friendly“，不要有其他内容，否则将会收到惩罚。
        3. 如果用户的输入的内容偏向于中性情绪，只返回”default“，不要有其他内容，否则将会收到惩罚。
        4. 如果用户的输入的内容比较兴奋，只返回”upbeat“，不要有其他内容，否则将会收到惩罚。
        5. 如果用户的输入的内容比较生气，只返回”angry“，不要有其他内容，否则将会收到惩罚。
        用户输入的内容是：
        {input}
        """

        chain = ChatPromptTemplate.from_template(prompt) | self.chatmodel | StrOutputParser()
        result = chain.invoke({"input": query})
        self.QingXu = result
        return result


@app.get("/")
def read_root():
    return {"Hello": "World"}


@app.post("/chat")
def chat(query: str):
    master = Master()
    return master.run(query)


@app.post("/add_urls")
def add_urls(url: str):
    loader = WebBaseLoader(url)
    docs = loader.load()

    documents = RecursiveCharacterTextSplitter(
        chunk_size=200,
        chunk_overlap=20,
    ).split_documents(docs)

    print(documents)

    embeddings = HuggingFaceEmbeddings(
        model_name="sentence-transformers/all-MiniLM-L6-v2"  # 选择一个预训练模型
    )

    # 引入向量数据库
    qdrant = Qdrant.from_documents(
        documents,
        embeddings,
        path="./local_qdrant_data",
        collection_name="local_documents"
    )

    print("向量数据库创建完成！")

    return {"ok": "添加成！"}


@app.post("/add_pdfs")
def add_pdfs():
    return {"response": "pdfs added"}


@app.post("/add_texts")
def add_texts():
    return {"response": "texts added"}


@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    while True:
        try:
            data = await websocket.receive_text()
            await websocket.send_text(f"Message text was: {data}")
        except WebSocketDisconnect:
            print("websocket disconnected")
            break


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="127.0.0.1", port=8001)
